#include "gemm.h"
#include "im2col.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
//#include <int8_t.h>
#include <string.h>
#include <stdint.h>
#ifdef _WIN32
#include <intrin.h>
#endif
// #if defined(_OPENMP)
// #include <omp.h>
// #endif

#define TILE_M 4 // 4 ops
#define TILE_N 16 // AVX2 = 2 ops * 8 int8_ts
#define TILE_K 16 // loop

#ifdef __cplusplus
#define PUT_IN_REGISTER
#else
#define PUT_IN_REGISTER register
#endif

// A(Weight) = M*la  B(im2col.T) = lb * K  C(out) = lc * M
void gemm_T(int TA, int TB, int M, int N, int K, int ALPHA,
    int8_t* A, int lda,
    int8_t* B, int ldb,
    int BETA,
    int8_t* C, int ldc,int shift)
{
    // printf("%d %d %d %d %d %d\n",M,N,K,lda,ldb,ldc);
    int i, j, k;
    for(i = 0; i<ldc; i++){
        for(j=0; j<M; j++){
            int temp = 0;
            for(k=0; k<K; k++){
                temp += A[j*lda+k] * B[i * K + k];
            }
            temp = temp>>shift;
            temp = temp>127? 127:temp;
            temp = temp<-127?-127:temp;
            C[i*M + j] = temp ; 
            // printf("%d ",temp);
        }
        // printf("\n");
    }

}

// void gemm_nt(int M, int N, int K, int8_t ALPHA,
//         int8_t *A, int lda,
//         int8_t *B, int ldb,
//         int8_t *C, int ldc)
// {
//     int i,j,k;
//     for(i = 0; i < M; ++i){
//         for(j = 0; j < N; ++j){
//             PUT_IN_REGISTER int8_t sum = 0;
//             for(k = 0; k < K; ++k){
//                 sum += ALPHA*A[i*lda+k]*B[j*ldb + k];
//             }
//             C[i*ldc+j] += sum;
//         }
//     }
// }

// void gemm_tn(int M, int N, int K, int8_t ALPHA,
//         int8_t *A, int lda,
//         int8_t *B, int ldb,
//         int8_t *C, int ldc)
// {
//     int i,j,k;
//     for(i = 0; i < M; ++i){
//         for(k = 0; k < K; ++k){
//             PUT_IN_REGISTER int8_t A_PART = ALPHA * A[k * lda + i];
//             for(j = 0; j < N; ++j){
//                 C[i*ldc+j] += A_PART*B[k*ldb+j];
//             }
//         }
//     }
// }

void gemm_nn(int M, int N, int K, int8_t ALPHA,
    int8_t* A, int lda,
    int8_t* B, int ldb,
    int8_t* C, int ldc)
{
    int i, j, k;
    for (i = 0; i < M; ++i) {
        for (k = 0; k < K; ++k) {
            PUT_IN_REGISTER int8_t A_PART = ALPHA * A[i * lda + k];
            //printf("A[%f] = %f, ALPHA:%f\n", i * lda + k, A[i * lda + k], ALPHA);
            for (j = 0; j < N; ++j) {
                C[i * ldc + j] += A_PART * B[k * ldb + j];
                //printf("C[%d] = %f, Apart:%f\n", i * ldc + j, C[i * ldc + j], A_PART);
            }
        }
    }
}

// void gemm_tt(int M, int N, int K, int8_t ALPHA,
//         int8_t *A, int lda,
//         int8_t *B, int ldb,
//         int8_t *C, int ldc)
// {
//     int i,j,k;
//     for(i = 0; i < M; ++i){
//         for(j = 0; j < N; ++j){
//             PUT_IN_REGISTER int8_t sum = 0;
//             for(k = 0; k < K; ++k){
//                 sum += ALPHA*A[i+k*lda]*B[k+j*ldb];
//             }
//             C[i*ldc+j] += sum;
//         }
//     }
// }


void gemm(int TA, int TB, int M, int N, int K, int8_t ALPHA,
    int8_t* A, int lda,
    int8_t* B, int ldb,
    int8_t BETA,
    int8_t* C, int ldc)
{
    gemm_cpu(TA, TB, M, N, K, ALPHA, A, lda, B, ldb, BETA, C, ldc);
}

// void gemm_nn_fast(int M, int N, int K, int8_t ALPHA,
//     int8_t* A, int lda,
//     int8_t* B, int ldb,
//     int8_t* C, int ldc)
// {
//     int i, j, k;
// //#pragma omp parallel for
//     for (i = 0; i < M; ++i) {
//         for (k = 0; k < K; ++k) {
//             PUT_IN_REGISTER int8_t A_PART = ALPHA * A[i * lda + k];
//             for (j = 0; j < N; ++j) {
//                 C[i * ldc + j] += A_PART * B[k * ldb + j];
//             }
//         }
//     }
// }

int is_avx() {
    return 0;
}

int is_fma_avx2() {
    return 0;
}

void gemm_cpu(int TA, int TB, int M, int N, int K, int8_t ALPHA,
    int8_t* A, int lda,
    int8_t* B, int ldb,
    int8_t BETA,
    int8_t* C, int ldc)
{
    //printf("cpu: %d %d %d %d %d %f %d %d %f %d\n",TA, TB, M, N, K, ALPHA, lda, ldb, BETA, ldc);
    if (BETA != 1) {
        int i, j;
        for (i = 0; i < M; ++i) {
            for (j = 0; j < N; ++j) {
                C[i * ldc + j] *= BETA;
                //printf("C[%d] = %d\n", i * ldc + j, C[i * ldc + j]);
            }
        }
    }
    gemm_nn(M, N, K, ALPHA, A, lda, B, ldb, C, ldc);
    // is_avx();   // initialize static variable
    // if (is_fma_avx2() && !TA && !TB) {
    //     gemm_nn_fast(M, N, K, ALPHA, A, lda, B, ldb, C, ldc);
    // }
    // else {
    //     int t;
    //     for (t = 0; t < M; ++t) {
    //         if (!TA && !TB)
    //             gemm_nn(1, N, K, ALPHA, A + t * lda, lda, B, ldb, C + t * ldc, ldc);
    //         else if (TA && !TB)
    //             gemm_tn(1, N, K, ALPHA, A + t, lda, B, ldb, C + t * ldc, ldc);
    //         else if (!TA && TB)
    //             gemm_nt(1, N, K, ALPHA, A + t * lda, lda, B, ldb, C + t * ldc, ldc);
    //         else
    //             gemm_tt(1, N, K, ALPHA, A + t, lda, B, ldb, C + t * ldc, ldc);
    //     }
    // }
}

